package com.dailywork.basepro.thread.concurrent;

import java.io.File;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.*;

/**
 * 计算文件夹大小
 *
 * @author 菠萝大象
 */
public class FileSizeCalc {

    static class SubDirsAndSize {

        public final long size;
        final List<File> subDirs;

        SubDirsAndSize(long size, List<File> subDirs) {
            this.size = size;
            this.subDirs = Collections.unmodifiableList(subDirs);
        }
    }

    private SubDirsAndSize getSubDirsAndSize(File file) {
        long total = 0;
        List<File> subDirs = new ArrayList<>();
        if (file.isDirectory()) {
            File[] children = file.listFiles();
            if (children != null) {
                for (File child : children) {
                    if (child.isFile()) {
                        total += child.length();
                    } else {
                        subDirs.add(child);
                    }
                }
            }
        }
        return new SubDirsAndSize(total, subDirs);
    }

    private long getFileSize(File file) throws Exception {
        final int cpuCore = Runtime.getRuntime().availableProcessors();
        final int poolSize = cpuCore + 1;
        ExecutorService service = Executors.newFixedThreadPool(poolSize);
        long total = 0;
        List<File> directories = new ArrayList<>();
        directories.add(file);
        SubDirsAndSize subDirsAndSize;
        try {
            while (!directories.isEmpty()) {
                List<Future<SubDirsAndSize>> partialResults = new ArrayList<>();
                for (final File directory : directories) {
                    //                    Callable<SubDirsAndSize> callable = () ->
                    // getSubDirsAndSize(directory);
                    //                    partialResults.add(service.submit(callable));
                    //                    Future<SubDirsAndSize> result = service.submit(() ->
                    // getSubDirsAndSize(directory));
                    //                    partialResults.add(result);

                    partialResults.add(
                            service.submit(
                                    new Callable<SubDirsAndSize>() {
                                        @Override
                                        public SubDirsAndSize call() throws Exception {
                                            return getSubDirsAndSize(directory);
                                        }
                                    }));
                }
                directories.clear();
                for (Future<SubDirsAndSize> partialResultFuture : partialResults) {
                    subDirsAndSize = partialResultFuture.get(100, TimeUnit.SECONDS);
                    total += subDirsAndSize.size;
                    directories.addAll(subDirsAndSize.subDirs);
                }
            }
            return total;
        } finally {
            service.shutdown();
        }
    }

    public static void main(String[] args) throws Exception {
        final long start = System.currentTimeMillis();
        long total = new FileSizeCalc().getFileSize(new File("H:\\gitbook"));
        final long end = System.currentTimeMillis();
        System.out.format("dir size: %dMB%n", total / (1024 * 1024));
        System.out.format("used time: %.3fs%n", (end - start) / 1.0e3);
    }
}
